module.exports = function ({ types: t }) {
  return {
    visitor: {
      Program(path) {
        let usesPropTypes = false;

        // Check if `propTypes` or `forbidExtraProps` is used anywhere in the file
        path.traverse({
          AssignmentExpression(assignPath) {
            if (
              t.isMemberExpression(assignPath.node.left) &&
              assignPath.node.left.property.name === "propTypes"
            ) {
              usesPropTypes = true;
            }
          },
        });

        if (usesPropTypes) {
          // Add `_checkPropTypes` import at the top of the file
          const importDeclaration = t.importDeclaration(
            [t.importDefaultSpecifier(t.identifier("_checkPropTypes"))],
            t.stringLiteral("prop-types/checkPropTypes")
          );
          path.unshiftContainer("body", importDeclaration);
        }
      },
      FunctionDeclaration(path) {
        insertPropTypesCheck(path, t);
      },
      VariableDeclarator(path) {
        // Handle arrow functions assigned to variables
        if (t.isArrowFunctionExpression(path.node.init)) {
          insertPropTypesCheck(path, t);
        }
      },
    },
  };

  function insertPropTypesCheck(path, t) {
    const node = path.node;
    let componentName;

    // Determine the component name
    if (t.isFunctionDeclaration(node)) {
      componentName = node.id.name;
    } else if (t.isVariableDeclarator(node)) {
      componentName = node.id.name;
    } else {
      return;
    }

    // Traverse the program body to find the `propTypes` assignment
    const programNode = path.findParent((parentPath) =>
      t.isProgram(parentPath.node)
    );

    if (!programNode || !programNode.node.body) {
      return; // Exit if the program body is not found
    }

    const programBody = programNode.node.body;

    const propTypesAssignment = programBody.find(
      (n) =>
        t.isExpressionStatement(n) &&
        t.isAssignmentExpression(n.expression) &&
        t.isMemberExpression(n.expression.left) &&
        n.expression.left.object.name === componentName &&
        n.expression.left.property.name === "propTypes"
    );

    if (!propTypesAssignment) {
      return; // Skip if no propTypes assignment found
    }

    let propTypesNode = propTypesAssignment.expression.right;

    // Check if propTypes is wrapped in forbidExtraProps
    if (
      t.isCallExpression(propTypesNode) &&
      t.isIdentifier(propTypesNode.callee) &&
      propTypesNode.callee.name === "forbidExtraProps"
    ) {
      // Insert the _checkPropTypes call with the wrapped propTypes
      const checkPropTypesCall = t.expressionStatement(
        t.callExpression(t.identifier("_checkPropTypes"), [
          t.cloneNode(propTypesNode), // Pass the wrapped propTypes (the result of forbidExtraProps)
          t.identifier("arguments[0]"),
          t.stringLiteral("prop"),
          t.stringLiteral(componentName),
        ])
      );

      // Insert `_checkPropTypes` call at the start of the component
      const functionBody =
        t.isFunctionDeclaration(node)
          ? node.body.body
          : node.init.body.body;

      functionBody.unshift(checkPropTypesCall);
    } else {
      // If not wrapped in forbidExtraProps, handle normally
      const checkPropTypesCall = t.expressionStatement(
        t.callExpression(t.identifier("_checkPropTypes"), [
          t.cloneNode(propTypesNode),
          t.identifier("arguments[0]"),
          t.stringLiteral("prop"),
          t.stringLiteral(componentName),
        ])
      );

      const functionBody =
        t.isFunctionDeclaration(node)
          ? node.body.body
          : node.init.body.body;

      functionBody.unshift(checkPropTypesCall);
    }
  }
};